{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2023-05-10T05:55:50.615727700Z",
     "start_time": "2023-05-10T05:55:50.508730400Z"
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.append('..')\n",
    "from common.time_layers import *\n",
    "from seq2seq import Encoder, Seq2seq\n",
    "from attention_layer import TimeAttention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "class AttentionEncoder(Encoder):\n",
    "    def forward(self, xs):\n",
    "        xs = self.embed.forward(xs)\n",
    "        hs = self.lstm.forward(xs)\n",
    "        return hs\n",
    "\n",
    "    def backward(self, dhs):\n",
    "        dout = self.lstm.backward(dhs)\n",
    "        dout = self.embed.backward(dout)\n",
    "        return dout"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-05-10T05:57:22.221371600Z",
     "start_time": "2023-05-10T05:57:22.214375500Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "class AttentionDecoder:\n",
    "    def __init__(self, vocab_size, wordvec_size, hidden_size):\n",
    "        V, D, H = vocab_size, wordvec_size, hidden_size\n",
    "        rn = np.random.randn\n",
    "\n",
    "        embed_W = (rn(V, D) / 100).astype('f')\n",
    "        lstm_Wx = (rn(D, 4 * H) / np.sqrt(D)).astype('f')\n",
    "        lstm_Wh = (rn(H, 4 * H) / np.sqrt(H)).astype('f')\n",
    "        lstm_b = np.zeros(4 * H).astype('f')\n",
    "        affine_W = (rn(2 * H, V) / np.sqrt(2 * H)).astype('f')\n",
    "        affine_b = np.zeros(V).astype('f')\n",
    "\n",
    "        self.embed = TimeEmbedding(embed_W)\n",
    "        self.lstm = TimeLSTM(lstm_Wx, lstm_Wh, lstm_b, stateful=True)\n",
    "        self.attention = TimeAttention()\n",
    "        self.affine = TimeAffine(affine_W, affine_b)\n",
    "        layers = [self.embed, self.lstm, self.attention, self.affine]\n",
    "\n",
    "        self.params, self.grads = [], []\n",
    "        for layer in layers:\n",
    "            self.params += layer.params\n",
    "            self.grads += layer.grads\n",
    "\n",
    "    def forward(self, xs, enc_hs):\n",
    "        h = enc_hs[:, -1]\n",
    "        self.lstm.set_state(h)\n",
    "\n",
    "        out = self.embed.forward(xs)\n",
    "        dec_hs = self.lstm.forward(out)\n",
    "        c = self.attention.forward(enc_hs, dec_hs)\n",
    "        out = np.concatenate((c, dec_hs), axis=2)\n",
    "        score = self.affine.forward(out)\n",
    "\n",
    "        return score\n",
    "\n",
    "    def backward(self, dscore):\n",
    "        dout = self.affine.backward(dscore)\n",
    "        N, T, H2 = dout.shape\n",
    "        H = H2 // 2\n",
    "\n",
    "        dc, ddec_hs0 = dout[:, :, :H], dout[:, :, H:]\n",
    "        denc_hs, ddec_hs1 = self.attention.backward(dc)\n",
    "        ddec_hs = ddec_hs0 + ddec_hs1\n",
    "        dout = self.lstm.backward(ddec_hs)\n",
    "        dh = self.lstm.dh\n",
    "        denc_hs[:, -1] += dh\n",
    "        self.embed.backward(dout)\n",
    "\n",
    "        return denc_hs\n",
    "\n",
    "    def generate(self, enc_hs, start_id, sample_size):\n",
    "        sampled = []\n",
    "        sample_id = start_id\n",
    "        h = enc_hs[:, -1]\n",
    "        self.lstm.set_state(h)\n",
    "\n",
    "        for _ in range(sample_size):\n",
    "            x = np.array([sample_id]).reshape((1, 1))\n",
    "\n",
    "            out = self.embed.forward(x)\n",
    "            dec_hs = self.lstm.forward(out)\n",
    "            c = self.attention.forward(enc_hs, dec_hs)\n",
    "            out = np.concatenate((c, dec_hs), axis=2)\n",
    "            score = self.affine.forward(out)\n",
    "\n",
    "            sample_id = np.argmax(score.flatten())\n",
    "            sampled.append(sample_id)\n",
    "\n",
    "        return sampled"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-05-10T05:58:23.370430900Z",
     "start_time": "2023-05-10T05:58:23.358430200Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "class AttentionSeq2seq(Seq2seq):\n",
    "    def __init__(self, vocab_size, wordvec_size, hidden_size):\n",
    "        args = vocab_size, wordvec_size, hidden_size\n",
    "        self.encoder = AttentionEncoder(*args)\n",
    "        self.decoder = AttentionDecoder(*args)\n",
    "        self.softmax = TimeSoftmaxWithLoss()\n",
    "        self.params = self.encoder.params + self.decoder.params\n",
    "        self.grads = self.encoder.grads + self.decoder.grads"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-05-10T05:59:28.120398400Z",
     "start_time": "2023-05-10T05:59:28.110399Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
